import numpy as np
import random
import matplotlib.pyplot as plt

def loadDataSet(filename):
    data_get = []
    fp = open(filename, 'r')
    for line in fp:
        curline = line.split(' ')
        floatline = list(map(float, curline))
        data_get.append(floatline)
    return data_get


def findCentroids(data_get, k):
    m = random.sample(data_get, k)
    return np.array(m)


def distEclud(vecA, vecB):
    return np.sqrt(sum(np.power(vecA - vecB, 2)))


def Kmeans(data_get, centroidsList):
    distculde = {}
    flag = 0
    for data in data_get:
        vecA = np.array(data)
        minDis = float('inf')
        for i in range(len(centroidsList)):
            vecB = centroidsList[i]
            distance = distEclud(vecA, vecB)
            if distance < minDis:
                minDis = distance
                flag = i
        if flag not in distculde.keys():
            distculde[flag] = list()
        distculde[flag].append(data)
    return distculde


def getCentroids(distculde):
    newcentroidsList = []
    for key in distculde:
        cent = np.array(distculde[key])
        newcentroid = np.mean(cent, axis=0)
        newcentroidsList.append(newcentroid.tolist())
    return np.array(newcentroidsList)


def calculate_Var(distculde, centroidsList):
    item_sum = 0.0
    for key in distculde:
        vecA = centroidsList[key]
        dist = 0.0
        for item in distculde[key]:
            vecB = np.array(item)
            dist += distEclud(vecA, vecB)
        item_sum += dist
    return item_sum


def showCluster(distculde, centroidsList):
    x = []
    y = []
    x.append(centroidsList[:, 0].tolist())
    y.append(centroidsList[:, 1].tolist())
    plt.plot(x, y, 'k*')
    colourList = ['bo', 'ro', 'yo', 'co']
    for i in distculde:
        centx = []
        centy = []
        for item in distculde[i]:
            centx.append(item[0])
            centy.append(item[1])
        plt.plot(centx, centy, colourList[i])
    plt.show()


if __name__ == "__main__":
    k = 4
    datMat = loadDataSet('1.txt')
    centroidsList = findCentroids(datMat, k)
    distculde = Kmeans(datMat, centroidsList)
    newVar = calculate_Var(distculde, centroidsList)
    oldVar = -0.0001
    while abs(newVar - oldVar) >= 0.0001:
        centroidsList = getCentroids(distculde)
        distculde = Kmeans(datMat, centroidsList)
        oldVar = newVar
        newVar = calculate_Var(distculde, centroidsList)
    showCluster(distculde, centroidsList)

